from __init__ import Predictor
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from domains import cartpole,acrobot
import random
import Tkinter
import math
import threading
import numpy
from itertools import izip


class CartPoleDrawer():
    def __init__(self,canvas):
        self.canvas = canvas
        
    def draw(self,state,step):
        if self.canvas is None: 
            return
        linelength=3
        rootx = state[0] + 2.5
        rooty = 5
        linex = math.sin(state[2])*linelength+rootx
        liney = rooty-math.cos(state[2])*linelength
        
        self.canvas.delete(Tkinter.ALL)
        self.canvas.create_line(rootx*100,rooty*100,linex*100,liney*100)
        self.canvas.create_line(5,5,10*(step % 2),10*(step % 2))
        
class AcrobotDrawer():
    def __init__(self,canvas):
        self.canvas = canvas
        
    def draw(self,state,step):
        if self.canvas is None: 
            return
        rootx = 100
        rooty = 100
        locs = acrobot.Acrobot.arm_kinematics(state[0], state[1])
        self.canvas.delete(Tkinter.ALL)
        self.canvas.create_line(rootx+0,rooty+2,rootx+locs[0],rooty-locs[1])
        self.canvas.create_line(rootx+locs[0],rooty-locs[1],rootx+locs[2],rooty-locs[3])
        self.canvas.create_line(5,5,10*(step % 2),10*(step % 2))
        
        
class Planner():
    EXPLORE_RATE = 0#0.05
    MAXDEPTH = 10
    def __init__(self, actionChoices, goalFunc,permitFunc):
        self.actionChoices = actionChoices
        self.plan = []
        self.goalFunc = goalFunc
        self.permitFunc = permitFunc
        self.turn = 0

    def reset(self):
        self.plan = []
        self.turn = 0
        
    @staticmethod
    def dist(pointA, pointB):
        try:
            d = 0.0
            for (a,b) in izip(pointA,pointB):
                if a is None or b is None: 
                    continue
                d +=  float(float(abs(a-b))**2.0)
        except:
            return float("inf")
        return d ** (0.5)
    
    @classmethod
    def states_close(cls, stateA,stateB,depth):
        dist = cls.dist(stateA, stateB)
        return dist < .01*(depth**0.5)
    
    def planner(self, state, actions, predictor,turn):
        #STATE state: currentstate
        #ITERATOR actions: available actions
        #FUNC predictor: (curState,action)->(nextState1,nextState2,nextState3...) 
        nextMove = None
        comesFrom = dict()
        stateQueue = [(state,0.0)]
        
        if turn % 100 == 0 and turn > 0:
            print turn

        while len(stateQueue) > 0:
            (currentState, curDepth) = stateQueue.pop(0)
            if curDepth > self.MAXDEPTH:
                # Should default to previously found path...
                break
            
            #if currentComp is similar to currentState, break out and choose that action.

            isnear = self.goalFunc(currentState,curDepth)
            if isnear:
                self.plan = []
                if curDepth == 0:
                    break
                s = currentState
                action = None
                while s in comesFrom:
                    action = comesFrom[s][1]
                    self.plan.append(action)
                    
                    s = comesFrom[s][0]
                     
                nextMove = self.plan.pop()
                
                break
            for action in actions:
                for nextState in predictor(currentState, action):                        
                    if (not self.permitFunc(nextState)) and self.permitFunc(state):
                        #ignore unacceptable states
                        continue
                    if nextState in comesFrom: 
                        #ignore repeats that we can get to from another (potentially shorter) path.
                        continue
                    stateQueue.append((nextState, curDepth+1))
                    comesFrom[nextState] = (currentState,action) 

        if nextMove is not None:
            if random.random() < self.EXPLORE_RATE:
                nextMove = None
        else:
            if len(self.plan) > 0:
                nextMove = self.plan.pop()
            else:
                nextMove = None
        return nextMove
        


def acrobotTest():
    def plotGNG(m):
        fig= plt.figure()
        
        for i,a, in enumerate((0,2)):
            ax = fig.add_subplot(1,2,i+1,projection='3d',zlim=(-.5,.5))
            for comp in m.gngs[a].graph.connected_components():
                gngnodes = [n.data.pos for n in comp]
                x = [n[0] for n in gngnodes]
                z = [n[0+4] for n in gngnodes]
                y = [n[1] for n in gngnodes]
                #ax.scatter(x,y,z)
                #pylab.matplotlib.pyplot.show()
                ax.scatter(x,y,z)
            
            
            x,y = numpy.meshgrid(numpy.arange(-.1,.1,.005),numpy.arange(-1,1,.2))
            zp,za = [],[]
        
            for yi in range(len(y)):
                zp.append([])
                za.append([])
                for xi in range(len(x[yi])):
                    sa = ((x[yi][xi],y[yi][xi],0,0),a)
                    pred = tuple(m.possActionResults(*sa))
                    if len(pred) == 0:
                        p = 0
                    else:
                        p = pred[-1][0]
                    zp[-1].append(p)
                    za[-1].append(acrobot.Acrobot.nextState(*sa)[0])
            
            ax.plot_wireframe(x,y,zp,color='b')
            ax.plot_wireframe(x,y,za,color='r')
        plt.show()  
        
    RANDOMTRAIN = False
    Master = Tkinter.Tk()
    w = Tkinter.Canvas(Master, width=600, height=600)
        
    TRAINCOUNT = 200
    acrodraw = AcrobotDrawer(w)
    p = Planner((0,1,2),
                lambda s, d:Planner.states_close(s, (math.pi,0,None,None), d),
                lambda s: True
                )
    
    m = Predictor(p,(0,1,2),4,
                acrodraw)
    epcount = 0
    best = 0
    results = []
    w.pack()
    th = threading.Thread(target=Master.mainloop)
    th.start()
    while epcount < TRAINCOUNT:
        ac = acrobot.Acrobot(math.pi,0,0,0)                
        if RANDOMTRAIN:
            t = ac.single_episode(m.testDrive)
            
        else:
            t = ac.single_episode(m.smartDrive)
    
        if len(t) < 1000:
            m.reset() 
    
        print epcount, len(t)
        results.append(len(t))
        best = max(len(t),best)
        epcount += 1
    
        plotGNG(m)
            
    m.save()
    print "Best: %s" % (best,)
    
    Master.destroy()
    th.join()
    
    import pylab
    
    if len(results) > 0:
        pylab.scatter(*zip(*enumerate(results)))
        pylab.show()

def cartPoleTest():
    def plotGNG(m):
        fig= plt.figure()
    
        for i,a, in enumerate((0,1)):
            ax = fig.add_subplot(1,2,i+1,projection='3d',zlim=(-.5,.5))
            for comp in m.gngs[a].graph.connected_components():
                gngnodes = [n.data.pos for n in comp]
                x = [n[2] for n in gngnodes]
                z = [n[2+4] for n in gngnodes]
                y = [n[3] for n in gngnodes]
                #ax.scatter(x,y,z)
                #pylab.matplotlib.pyplot.show()
                ax.scatter(x,y,z)
            
            
            x,y = numpy.meshgrid(numpy.arange(-.5,.52,.05),numpy.arange(-3,3,.5))
            zp,za = [],[]
        
            for yi in range(len(y)):
                zp.append([])
                za.append([])
                for xi in range(len(x[yi])):
                    sa = ((0,0,x[yi][xi],y[yi][xi]),a)
                    pred = tuple(m.possActionResults(*sa))
                    if len(pred) == 0:
                        p = 0
                    else:
                        p = pred[-1][2]
                    zp[-1].append(p)
                    za[-1].append(cartpole.CartPole.nextState(*sa)[2])
            
            ax.plot_wireframe(x,y,zp,color='b')
            ax.plot_wireframe(x,y,za,color='r')
        plt.show() 
        
         
    RANDOMTRAIN = False
    Master = Tkinter.Tk()
    w = Tkinter.Canvas(Master, width=600, height=600)
    cartdraw = CartPoleDrawer(w)
    w.pack()
    th = threading.Thread(target=Master.mainloop)
    th.start()
    
    random.seed(3)
    TRAINCOUNT = 200
    p = Planner((0,1),
                lambda s,d:Planner.states_close(s, (None,None,0,None), d),
                lambda s: abs(s[2]) < 0.2 and abs(s[0] < 2.0))
    
    cp = cartpole.CartPole()
    m = Predictor(p,(0,1),4,
                cartdraw)
    epcount = 0
    best = 0
    results = []
    
    while epcount < TRAINCOUNT:           
        if RANDOMTRAIN:
            t = cp.single_episode(m.testDrive)
            
        else:
            t = cp.single_episode(m.smartDrive)
    
        m.reset()
    
        print epcount, len(t)
        results.append(len(t))
        best = max(len(t),best)
        epcount += 1
        
        
    m.save()
    print "Best: %s" % (best,)
    
    
    Master.destroy()
    th.join()
    plotGNG(m)
    
    import pylab
    
    if len(results) > 0:
        pylab.scatter(*zip(*enumerate(results)))
        pylab.show()
 
cartPoleTest()        
#acrobotTest()